# -*- coding: utf-8 -*-
# @author: HRUN
import json
import subprocess

from locust import FastHttpUser, task, TaskSet, between, events, LoadTestShape, HttpUser, constant
from locust.env import Environment
from gevent._semaphore import Semaphore

# 集合点逻辑实现
all_locusts_spawned = Semaphore()
all_locusts_spawned.acquire()  # 阻塞线程
def on_hatch_complete(user_count, **kwargs):
    """
    Select_task类的钩子方法
    """
    print(f"{user_count}个虚拟用户产生完成")
    # 创建钩子方法
    all_locusts_spawned.release()
events.spawning_complete.add_listener(on_hatch_complete)



# 不同场景存在不同环境
class SubEnvironment(Environment):
    def __init__(self):
        super().__init__()
        self.url = None




class BaseStepDispose():
    def api_perform(self, step):
        data = step.get('content')
        path = data.get('url')
        params = self.construction_data(step)
        headers = self.construction_headers(step)
        if data.get('method') == 'POST':
            # TODO:场景环境需要自定义
            with self.client.post(f"{path}", data=json.dumps(params), headers=headers,
                catch_response=True) as res:
                    if res.status_code == 200:
                        print("入参：%s, 返回值:%s" % (params, res.text))

        else:
            return ''
            # TODO: 处理GET类型
        return self

    def construction_data(self, step):
        pass

    def construction_headers(self, step):
        pass


    def if_perform(self, case):
        pass

    def script_perform(self, case):
        pass




class GenerateTask():
    def __init__(self, control='20'):
        self.control = control
    def create_taskSet_class(self, item: list) -> list:
            task_class = []
            for idx, scence in enumerate(item):
                cls_name = scence.get('name') or 'Demo' + f'_{idx + 1}'
                env = scence.get('env', None)
                weight = int(scence.get('weight', 1))
                task_funcs = {}
                if scence.get('steps'):
                    for step_idx, step in enumerate(scence['steps']):
                        func_name = self.create_testTask_name(idx, step_idx, step)
                        test_method = self.create_testTask_func(step)
                        task_funcs[func_name] = test_method
                task_funcs['on_start'] = self.on_start
                task_funcs['on_stop'] = self.on_stop
                # TODO: 添加场景权重以及运行环境属性
                task_funcs['weight'] = weight
                task_funcs['host'] = 'https://baidu.com'
                cls = type(cls_name, (TaskSet,BaseStepDispose,), task_funcs)
                print(cls.__dict__)
                task_class.append(cls)

            return task_class

    def on_start(self):
        '''场景集初始化执行方法'''
        global n
        if self.control == '10':
            n += 1
            print("%s个虚拟用户开始启动" % n)
            all_locusts_spawned.wait()  # 同步锁等待
        print('开始执行场景集初始化动作')
    def on_stop(self):
        """场景集结束执行方法"""
        print('结束执行场景集初始化动作')
    def create_testTask_name(self, case_index, step_index, step):
        """生成唯一的测试方法名，避免重复"""
        step_title = step.get('title', 'UnknownTitle').replace(' ', '_')
        testTask_name = f"testTask_{case_index + 1}_{step_index + 1}_{step_title}"
        return testTask_name

    def create_testTask_func(self, step):
        """创建测试方法，使用函数工厂来动态传递参数"""
        type = step.get('type')
        def test_method(self):
            if type == 'api':
                return self.api_perform(step)
            elif type == 'script':
                return self.script_perform(step)
            elif type == 'if':
                return self.if_perform(step)
            else:
                raise ValueError(f"不支持的步骤类型:{type}")
        if type == 'api':
            weight = step.get('weight', 1)
            return task(weight)(test_method)

        return test_method


class Config():
    def __init__(self, config):
        self.config = config
        self.write_config()

    def write_config(self):
        """写入配置文件"""
        # config_lines = ["headless = true\n"]
        config_lines = ["headless = false\n"]
        if isinstance(self.config, dict):
            time_unit = None
            pressureMode = None
            for key, value in self.config.items():
                if key == 'timeUnit':
                    time_unit = value
                elif key == 'pressureMode':
                    pressureMode = value
                elif key == 'pressureConfig':
                    if pressureMode == '10':
                        config_lines.extend(self.process_pressure_config(value, time_unit))

                elif key == 'serverArray':
                    self.server(value)

        # 写入配置文件
        try:
            with open("locust.conf", "w") as config_file:
                config_file.writelines(config_lines)
        except IOError as e:
            print(f"Error writing to file: {e}")
            # 错误处理重试处理

    def process_pressure_config(self, pressure_config, time_unit) -> list:
        """处理 pressureConfig 配置"""
        lines = []
        for key, value in pressure_config.items():
            if key == 'lastLong' and time_unit:
                lines.append(f"run-time = {value}{time_unit}\n")
            elif key == 'concurrencyNumber':
                lines.append(f"users = {value}\n")
            elif key == 'concurrencyStep':
                lines.append(f"spawn-rate = {value}\n")

        return lines

    def server(self, server_array: list):
        """启动locust的服务器环境"""
        pass

    def config_file(self) -> dict:
        """配置项返回给其他模块使用"""
        config = self.config
        filter_config = {key: value for key, value in config.items() if key not in
                         ['id', 'serverNames','creator','create_time','modifier',
                          'update_time', 'name',  'isSetting', 'project',
                          'task','serverArray', 'resource', ]}

        stages = []
        if filter_config.get('pressureMode', None) == '20':
            config = filter_config.get('pressureConfig')
            ladders = config.get('ladders')
            sum_duration = 0
            for stage in ladders:
                lastLong = stage.get('lastLong',0)
                duration = self.timeUnit(filter_config.get('timeUnit', 's'), lastLong)
                sum_duration += duration
                stages.append({
                    "duration": sum_duration,
                    "users": int(stage.get('concurrencyNumber',0)),
                    "spawn_rate": int(stage.get('concurrencyStep', 0)),
                })
            filter_config['stages'] = stages
            del filter_config['pressureConfig']

        print(filter_config)
        return  filter_config


    def timeUnit(self, time_unit, value):
        """时间单位转换"""
        value = int(value)
        if time_unit == 's':
            return value
        elif time_unit == 'm':
            return value * 60
        elif time_unit == 'h':
            return value * 3600
        else:
            return value







item = {
  "scenes": [
    {
      "steps": [
        {
          "id": 254,
          "title": "HTTP接口",
          "content": {
            "id": 1,
            "name": "访问百度",
            "host": {},
            "url": "https://tmstest.gree.com/",
            "method": "POST",
            "headers": {},
            "request": {
              "json": {},
              "data": None,
              "params": {}
            },
            "file": [],
            "setup_script": "# 前置脚本(python):\n# global_tools:全局工具函数\n# data:用例数据 \n# env: 局部环境\n# ENV: 全局环境\n# db: 数据库操作对象",
            "teardown_script": "# 后置脚本(python):\n# global_tools:全局工具函数\n# data:用例数据 \n# response:响应对象response \n# env: 局部环境\n# ENV: 全局环境\n# db: 数据库操作对象",
            "interface_tag": {
              "tag": []
            },
            "creator": "admin",
            "create_time": "2024-07-12T18:47:59.044109+08:00",
            "modifier": "admin",
            "update_time": "2024-11-07T23:54:40.378000+08:00",
            "desc": "",
            "type": "api",
            "YApi_id": None,
            "YApi_status": 0,
            "status": "开发中",
            "project": 1,
            "treenode": 1
          },
          "type": "api",
          "script": "",
          "status": True,
          "weight": 1,
          "children": []
        },
        {
          "id": 261,
          "title": "定时控制器",
          "content": {
            "time": 0
          },
          "type": "time",
          "script": "",
          "status": True,
          "weight": None,
          "children": []
        }
      ],
      "name": "下单场景",
      "env": 1,
      "weight": 1
    },
    {
      "steps": [
        {
          "id": 253,
          "title": "HTTP接口",
          "content": {
            "id": 1,
            "name": "访问百度",
            "host": {},
            "url": "https://www.baidu.com/",
            "method": "POST",
            "headers": {},
            "request": {
              "json": {},
              "data": None,
              "params": {}
            },
            "file": [],
            "setup_script": "# 前置脚本(python):\n# global_tools:全局工具函数\n# data:用例数据 \n# env: 局部环境\n# ENV: 全局环境\n# db: 数据库操作对象",
            "teardown_script": "# 后置脚本(python):\n# global_tools:全局工具函数\n# data:用例数据 \n# response:响应对象response \n# env: 局部环境\n# ENV: 全局环境\n# db: 数据库操作对象",
            "interface_tag": {
              "tag": []
            },
            "creator": "admin",
            "create_time": "2024-07-12T18:47:59.044109+08:00",
            "modifier": "admin",
            "update_time": "2024-11-07T23:54:40.378000+08:00",
            "desc": "",
            "type": "api",
            "YApi_id": None,
            "YApi_status": 0,
            "status": "开发中",
            "project": 1,
            "treenode": 1
          },
          "type": "api",
          "script": "",
          "status": None,
          "weight": 1,
          "children": []
        }
      ],
      "name": "查询场景",
      "env": 4,
      "weight": 10
    }
  ],
  # "presetting": {
  #   "id": 51,
  #   "serverNames": ["毓燊的服务器", "系统管理"],
  #   "creator": "admin",
  #   "create_time": "2024-10-21T22:26:50.611933+08:00",
  #   "modifier": "admin",
  #   "update_time": "2024-11-14T21:50:41.046000+08:00",
  #   "name": "吃吃吃啊啊啊",
  #   "rule": "* * * * *",
  #   "taskType": "10",
  #   "logMode": "0",
  #   "pressureMode": "20",
  #   "timeUnit": "s",
  #   "control": "20",
  #   "resource": "20",
  #   "pressureConfig": {
  #     "lastLong": "32",
  #     "concurrencyNumber": "23",
  #     "thinkTime": "32"
  #   },
  #   "isSetting": True,
  #   "project": 1,
  #   "task": 11,
  #   "serverArray": [22, 25]
  # },
  "presetting": {
        "id": 73,
        "serverNames": [
            "yzy\u7684\u670d\u52a1\u5668",
            "\u6bd3\u71ca\u7684\u670d\u52a1\u5668"
        ],
        "creator": "admin",
        "create_time": "2024-11-28T14:33:41.535507+08:00",
        "modifier": "admin",
        "update_time": "2024-11-28T14:45:07.580000+08:00",
        "name": "\u914d\u7f6e2",
        "rule": "* * * * *",
        "taskType": "10",
        "logMode": "0",
        "pressureMode": "10",
        "timeUnit": "s",
        "control": "20",
        "resource": "20",
        "pressureConfig": {
            "lastLong": "32",
            "concurrencyNumber": "23",
            "concurrencyStep": "32"
        },
        "isSetting": True,
        "thinkTime": [
            1,
            3
        ],
        "thinkTimeType": "20",
        "project": 1,
        "task": 16,
        "serverArray": [
            21,
            22
        ]
    }
}

conf = Config(item.get('presetting', {}))
runConf = conf.config_file()
# 创建 GenerateTask 实例
generator = GenerateTask(runConf.get('control'))

cls_list = generator.create_taskSet_class(item.get('scenes'))



class CreateKitClass(FastHttpUser):
    tasks = cls_list
    host = ""
    if runConf.get('thinkTimeType') == '10':
        wait_time = constant(runConf.get('thinkTime')[0])
    else:
        wait_time = between(runConf.get('thinkTime')[0], runConf.get('thinkTime')[1])

    @events.test_start.add_listener
    def on_test_start(environment, **kwargs):
        print("开始执行任务集初始化动作")
    @events.test_stop.add_listener
    def on_test_stop(environment, **kwargs):
        print("结束执行任务集初始化动作")


if runConf.get('pressureMode') == '20':
    class StagesShapeWithCustomUsers(LoadTestShape):
        """
        A simply load test shape class that has different user and spawn_rate at
        different stages.

        Keyword arguments:
            stages -- A list of dicts, each representing a stage with the following keys:
                duration -- When this many seconds pass the test is advanced to the next stage #持续时间
                users -- Total user count 用户数
                spawn_rate -- Number of users to start/stop per second 每秒产生或停止的用户数
                user_classes -- 指定的任务让负载更精确的控制该任务，如果不指定就是随机的
                stop -- A boolean that can stop that test at a specific stage 要想在哪个阶段停止运行就设置该值

            stop_at_end -- Can be set to stop once all stages have run.
        """

        stages = runConf.get('stages',[])

        def tick(self):
            run_time = self.get_run_time()

            for stage in self.stages:
                if run_time < stage["duration"]:
                    # Not the smartest solution, TODO: find something better
                    try:
                        tick_data = (stage["users"], stage["spawn_rate"], stage["user_classes"])
                    except KeyError:
                        tick_data = (stage["users"], stage["spawn_rate"])
                    return tick_data

            return None

if __name__ == '__main__':
    subprocess.run(
        ["locust", "-f", "taskGenerate.py"])

